// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software distributed under the License
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied. See the License for the specific language governing permissions and limitations under the License

/**
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */

#include <aws/core/internal/AWSHttpResourceClient.h>
#include <aws/core/client/DefaultRetryStrategy.h>
#include <aws/core/http/HttpClient.h>
#include <aws/core/http/HttpClientFactory.h>
#include <aws/core/http/HttpResponse.h>
#include <aws/core/utils/logging/LogMacros.h>
#include <aws/core/utils/StringUtils.h>
#include <aws/core/utils/HashingUtils.h>
#include <aws/core/platform/Environment.h>
#include <aws/core/client/AWSError.h>
#include <aws/core/client/CoreErrors.h>
#include <aws/core/utils/xml/XmlSerializer.h>
#include <limits.h>
#include <mutex>
#include <sstream>
#include <random>
#include "milvus-storage/filesystem/s3/AliyunSTSClient.h"

namespace Aws {
namespace Http {
class HttpClient;
class HttpRequest;
enum class HttpResponseCode;
}  // namespace Http

namespace Client {
Aws::String ComputeUserAgentString();
}

namespace Internal {

static const char STS_RESOURCE_CLIENT_LOG_TAG[] = "AliyunSTSResourceClient";  // [aliyun]

int IntRand(const int& min, const int& max) {
  static thread_local std::mt19937 generator;
  std::uniform_int_distribution<int> distribution(min, max);
  return distribution(generator);
}

AliyunSTSCredentialsClient::AliyunSTSCredentialsClient(const Aws::Client::ClientConfiguration& clientConfiguration)
    : AWSHttpResourceClient(clientConfiguration, STS_RESOURCE_CLIENT_LOG_TAG) {
  m_aliyunOidcProviderArn = Aws::Environment::GetEnv("ALIBABA_CLOUD_OIDC_PROVIDER_ARN");
  if (m_aliyunOidcProviderArn.empty()) {
    AWS_LOGSTREAM_WARN(STS_RESOURCE_CLIENT_LOG_TAG,
                       "oidc role arn must be specified to use STS "
                       "AssumeRole web identity creds provider.");
    return;
  }
  SetErrorMarshaller(Aws::MakeUnique<Aws::Client::XmlErrorMarshaller>(STS_RESOURCE_CLIENT_LOG_TAG));

  // [aliyun]
  m_endpoint = "https://sts.aliyuncs.com";

  AWS_LOGSTREAM_INFO(STS_RESOURCE_CLIENT_LOG_TAG, "Creating STS ResourceClient with endpoint: " << m_endpoint);
}

AliyunSTSCredentialsClient::STSAssumeRoleWithWebIdentityResult
AliyunSTSCredentialsClient::GetAssumeRoleWithWebIdentityCredentials(
    const STSAssumeRoleWithWebIdentityRequest& request) {
  // Calculate query string
  Aws::StringStream ss;
  // [aliyun]
  // linux curl example:
  // Action=AssumeRoleWithOIDC
  // Timestamp=`date -Iseconds`
  // Version="2015-04-01"
  // SignatureNonce=`$RANDOM`
  // RoleSessionName="default_session"
  // RoleArn=$ALIBABA_CLOUD_ROLE_ARN
  // OIDCProviderArn=$ALIBABA_CLOUD_OIDC_PROVIDER_ARN
  // OIDCToken=`cat $ALIBABA_CLOUD_OIDC_TOKEN_FILE`
  // curl "https://sts.aliyuncs.com?Action=$Action&Timestamp=$time" \
            //     -H "Host: sts.aliyuncs.com" \
            //     -H "Accept-Encoding: identity" \
            //     -H "SignatureNonce: $SignatureNonce" \
            //     -d
  //     "RoleArn=$RoleArn&OIDCProviderArn=$OIDCProviderArn&OIDCToken=$OIDCToken&RoleSessionName=$RoleSessionName&Version=$Version"
  ss << "Action=AssumeRoleWithOIDC"
     << "&Timestamp=" /*iso8601*/
     << Aws::Utils::StringUtils::URLEncode(
            Aws::Utils::DateTime::Now().ToGmtString(Aws::Utils::DateFormat::ISO_8601).c_str())
     << "&Version=2015-04-01"
     << "&SignatureNonce="
     << Aws::Utils::HashingUtils::HashString(Aws::Utils::StringUtils::to_string(IntRand(0, INT_MAX)).c_str())
     << "&RoleSessionName=" << Aws::Utils::StringUtils::URLEncode(request.roleSessionName.c_str())
     << "&RoleArn=" << Aws::Utils::StringUtils::URLEncode(request.roleArn.c_str())
     << "&OIDCProviderArn=" << Aws::Utils::StringUtils::URLEncode(m_aliyunOidcProviderArn.c_str())
     << "&OIDCToken=" << Aws::Utils::StringUtils::URLEncode(request.webIdentityToken.c_str());

  std::shared_ptr<Aws::Http::HttpRequest> httpRequest(Aws::Http::CreateHttpRequest(
      m_endpoint, Aws::Http::HttpMethod::HTTP_POST, Aws::Utils::Stream::DefaultResponseStreamFactoryMethod));

  httpRequest->SetUserAgent(Aws::Client::ComputeUserAgentString());

  std::shared_ptr<Aws::IOStream> body = Aws::MakeShared<Aws::StringStream>("STS_RESOURCE_CLIENT_LOG_TAG");
  *body << ss.str();

  httpRequest->AddContentBody(body);
  body->seekg(0, body->end);
  auto streamSize = body->tellg();
  body->seekg(0, body->beg);
  Aws::StringStream contentLength;
  contentLength << streamSize;
  httpRequest->SetContentLength(contentLength.str());
  httpRequest->SetContentType("application/x-www-form-urlencoded");

  Aws::String credentialsStr = GetResourceWithAWSWebServiceResult(httpRequest).GetPayload();

  // Parse credentials
  STSAssumeRoleWithWebIdentityResult result;
  if (credentialsStr.empty()) {
    AWS_LOGSTREAM_WARN(STS_RESOURCE_CLIENT_LOG_TAG, "Get an empty credential from sts");
    return result;
  }

  // [aliyun] output example
  // <?xml version='1.0' encoding='UTF-8'?>
  // <AssumeRoleWithOIDCResponse>
  //     <RequestId>D94AFCC3-54CA-508C-B6AF-21481E761BDB</RequestId>
  //     <OIDCTokenInfo>
  //         <Issuer>https://oidc-ack-cn-shanghai.oss-cn-shanghai.aliyuncs.com/c532c4ce5e84048a1972535df283f737d</Issuer>
  //         <Subject>system:serviceaccount:default:my-release</Subject>
  //         <ClientIds>sts.aliyuncs.com</ClientIds>
  //     </OIDCTokenInfo>
  //     <AssumedRoleUser>
  //         <Arn>acs:ram::1413891078881348:role/vdc-poc-milvus/default</Arn>
  //         <AssumedRoleId>383373758575348335:default</AssumedRoleId>
  //     </AssumedRoleUser>
  //     <Credentials>
  //         <SecurityToken>xxx</SecurityToken>
  //         <AccessKeyId>xxx</AccessKeyId>
  //         <AccessKeySecret>xxx</AccessKeySecret>
  //         <Expiration>2023-03-02T07:39:09Z</Expiration>
  //     </Credentials>
  // </AssumeRoleWithOIDCResponse>
  const Utils::Xml::XmlDocument xmlDocument = Utils::Xml::XmlDocument::CreateFromXmlString(credentialsStr);
  Utils::Xml::XmlNode rootNode = xmlDocument.GetRootElement();
  Utils::Xml::XmlNode resultNode = rootNode;
  if (!rootNode.IsNull() && (rootNode.GetName() != "AssumeRoleWithOIDCResponse")) {
    resultNode = rootNode.FirstChild("AssumeRoleWithOIDCResponse");  // [aliyun]
  }

  if (!resultNode.IsNull()) {
    Utils::Xml::XmlNode credentialsNode = resultNode.FirstChild("Credentials");
    if (!credentialsNode.IsNull()) {
      Utils::Xml::XmlNode accessKeyIdNode = credentialsNode.FirstChild("AccessKeyId");
      if (!accessKeyIdNode.IsNull()) {
        result.creds.SetAWSAccessKeyId(accessKeyIdNode.GetText());
      }

      Utils::Xml::XmlNode secretAccessKeyNode = credentialsNode.FirstChild("AccessKeySecret");  // [aliyun]
      if (!secretAccessKeyNode.IsNull()) {
        result.creds.SetAWSSecretKey(secretAccessKeyNode.GetText());
      }

      Utils::Xml::XmlNode sessionTokenNode = credentialsNode.FirstChild("SecurityToken");  // [aliyun]
      if (!sessionTokenNode.IsNull()) {
        result.creds.SetSessionToken(sessionTokenNode.GetText());
      }

      Utils::Xml::XmlNode expirationNode = credentialsNode.FirstChild("Expiration");
      if (!expirationNode.IsNull()) {
        result.creds.SetExpiration(Aws::Utils::DateTime(
            Aws::Utils::StringUtils::Trim(expirationNode.GetText().c_str()).c_str(), Aws::Utils::DateFormat::ISO_8601));
      }
    }
  }
  return result;
}
}  // namespace Internal
}  // namespace Aws
